import pandas as pd
import spacy
from tqdm import tqdm


def spacy_ner(
    fp_input: str,
    fp_output: str, 
    model: str = "en_core_web_trf",
):

    qrys = pd.read_csv(fp_input, usecols=["qry", "qry_id"])
    qrys = qrys.drop_duplicates(subset="qry_id")

    nlp = spacy.load(model)
    all_ents = []
    for qry in tqdm(qrys.qry):
        ents = nlp(qry).ents
        if len(ents):
            all_ents.append([e.label_ for e in ents])
        else:
            all_ents.append("nonentity")

    qrys["type_trf"] = all_ents
    qrys = qrys.explode("type_trf")
    qrys.to_csv(fp_output, index=False)


if __name__ == "__main__":
    spacy_ner()
